# langchain_agent_demo_juhe_weather.py：对接聚合数据天气API的Agent代码
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import ZeroShotAgent, AgentExecutor
from langchain.prompts import PromptTemplate
# 修复记忆组件导入（抑制警告）
import warnings
from langchain.memory import ConversationBufferMemory

warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*ConversationBufferMemory.*")

from dotenv import load_dotenv
import os

# 1. 加载环境变量（聚合API密钥可写在.env，也可直接使用提供的密钥）
load_dotenv()


# -------------------------- 工具1：天气查询（替换为聚合数据API） --------------------------
@tool(
    description="用于查询指定城市的当前实时天气。参数：city（必填，中文城市名，示例：苏州、北京、上海），返回值：实时温度、湿度、天气状况、风向风力。调用聚合数据天气API获取数据，不可编造。")
def WeatherTool(city: str) -> str:
    # 聚合数据天气API配置（固定参数）
    api_url = "http://apis.juhe.cn/simpleWeather/query"
    api_key = "89d99d8a8f8f20938ebfb99911c63367"  # 提供的固定密钥
    # 若想从.env文件读取（更安全），可替换为：api_key = os.getenv("JUHE_WEATHER_API_KEY")

    # 构造请求参数（聚合API支持中文城市名，无需编码）
    request_params = {
        "key": api_key,
        "city": city.strip()  # 去除空格，避免参数错误
    }

    try:
        import requests
        # 发送GET请求（聚合API支持GET/POST，GET更简单）
        response = requests.get(
            url=api_url,
            params=request_params,
            timeout=8  # 超时时间8秒，避免卡壳
        )
        response.raise_for_status()  # 若HTTP状态码非200，抛出异常（如404、500）

        # 解析JSON响应
        result_data = response.json()
        error_code = result_data.get("error_code")

        # 1. 判断API请求是否成功（error_code=0为成功）
        if error_code != 0:
            reason = result_data.get("reason", "未知错误")
            return f"获取{city}天气失败：API返回错误（{error_code}）- {reason}"

        # 2. 提取实时天气数据（从result.realtime字段）
        realtime_weather = result_data.get("result", {}).get("realtime")
        if not realtime_weather:
            return f"获取{city}天气失败：API返回数据中无实时天气信息"

        # 3. 整理需要的字段（温度、湿度、天气状况、风向、风力）
        temperature = realtime_weather.get("temperature", "未知")  # 温度（如"33"）
        humidity = realtime_weather.get("humidity", "未知")  # 湿度（如"58"）
        weather_info = realtime_weather.get("info", "未知")  # 天气状况（如"多云"）
        direct = realtime_weather.get("direct", "未知")  # 风向（如"西风"）
        power = realtime_weather.get("power", "未知")  # 风力（如"1级"）

        # 4. 格式化返回结果（清晰易读）
        return (
            f"{city}当前实时天气：\n"
            f"- 天气状况：{weather_info}\n"
            f"- 温度：{temperature}℃\n"  # 补充℃符号，更直观
            f"- 湿度：{humidity}%\n"  # 补充%符号
            f"- 风向：{direct}\n"
            f"- 风力：{power}"
        )

    # 捕获常见异常
    except requests.exceptions.Timeout:
        return f"获取{city}天气失败：请求超时（请检查网络是否正常）"
    except requests.exceptions.ConnectionError:
        return f"获取{city}天气失败：网络连接错误（无法访问API服务器）"
    except requests.exceptions.HTTPError as e:
        return f"获取{city}天气失败：HTTP错误（{str(e)}）"
    except Exception as e:
        return f"获取{city}天气失败：程序错误（{str(e)}）"


# -------------------------- 工具2：计算器（保持不变） --------------------------
@tool(
    description="用于计算加减乘除运算。参数：expression（必填，数学表达式，示例：2+3*4、10-5，无空格），返回值：计算结果。仅处理数字和+、-、*、/。")
def CalculatorTool(expression: str) -> str:
    try:
        allowed_chars = "0123456789+-*/."
        if not all(c in allowed_chars for c in expression):
            return "计算失败：表达式仅支持数字和 +、-、*、/、.（无空格）"
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except ZeroDivisionError:
        return "计算失败：不能除以0"
    except Exception as e:
        return f"计算失败：{str(e)}（示例正确格式：150+200+100）"


# -------------------------- 工具3：RAG查询（保持不变） --------------------------
@tool(
    description="用于查询公司内部知识库（如年假政策、产品价格）。参数：question（必填，查询问题，示例：公司年假政策、产品A价格），返回值：知识库答案。公司相关问题必须调用此工具，不可编造。")
def RAGTool(question: str) -> str:
    rag_data = {
        "公司年假政策": "员工入职满1年可享受5天年假，满3年10天，满5年15天。",
        "产品A的价格": "产品A基础版999元/年，高级版1999元/年，企业版需联系销售定制。",
        "加班补贴标准": "工作日加班每小时补贴20元，周末30元/小时，法定节假日50元/小时。"
    }
    for key, answer in rag_data.items():
        if key in question or any(kw in question for kw in key.split(" ")):
            return answer
    return f"知识库中未找到「{question}」的信息，请联系行政或产品部门获取详情。"


# 2. 整理工具列表
tools = [WeatherTool, CalculatorTool, RAGTool]

# 3. 初始化 LLM（通义千问 qwen-turbo）
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo",
    temperature=0,
    timeout=10
)

# 4. 提示词模板（保持优化后的记忆格式）
prompt_template = """
你是帮助用户解决问题的智能Agent，可使用以下工具：
{tools}
可用工具名称：{tool_names}

{chat_history}  # 历史对话（格式：用户：xxx\nAgent：xxx）

使用规则：
1. 必须先查看历史对话：若用户追问，优先从历史中提取关键数据（如之前的预算、天气）；
2. 调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数无空格）；
3. 涉及数字、天气等关键信息时，Final Answer必须明确强调；
4. 若工具调用失败，需明确告知用户原因，不要编造数据。

用户当前问题：{input}
思考过程：{agent_scratchpad}
"""

# 5. 创建提示词
prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "tool_names", "chat_history", "input", "agent_scratchpad"]
)

# 6. 创建 ZeroShotAgent
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True
)
"""
memory_key="chat_history"：指定在prompt中表示对话历史记录的键名，与PromptTemplate中的{chat_history}变量对应。
return_messages=False：表示以字符串形式返回对话历史，而不是以消息对象列表的形式返回。设为False时，历史记录会格式化为"用户：xxx\nAgent：xxx"的形式。
human_prefix="用户"：在对话历史中，用户说的话会以"用户："作为前缀。
ai_prefix="Agent"：在对话历史中，AI的回答会以"Agent："作为前缀。
"""

# 7. 初始化记忆组件
memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=False,
    human_prefix="用户",
    ai_prefix="Agent"
)

# 8. 创建 Agent 执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
    memory=memory
)

# 9. 测试：查询苏州天气（验证聚合API）+ 预算追问
if __name__ == "__main__":
    # 场景1：查询苏州实时天气（验证新接口）
    print("=" * 50)
    print("场景1：查询苏州实时天气（聚合API）")
    # 查看对话历史
    print("对话历史:")
    print(memory.buffer)
    result1 = agent_executor.invoke({"input": "查询苏州今天的实时天气。"})
    print("Agent 回答：", result1["output"])
    print("\n" + "-" * 30 + "\n")

    # 场景2：多工具调用（上海天气+穿搭+预算）
    print("场景2：苏州天气+穿搭+预算")
    # 查看对话历史
    print("对话历史:")
    print(memory.buffer)
    result2 = agent_executor.invoke({"input": "查苏州今天的实时天气，推荐穿搭，再算预算（T恤180元，短裤120元）。"})
    print("Agent 回答：", result2["output"])
    print("\n" + "-" * 30 + "\n")

    # 场景3：追问加预算（测试记忆）
    print("场景3：追问加预算（测试记忆）")
    # 查看对话历史
    print("对话历史:")
    print(memory.buffer)
    print(memory.load_memory_variables({}))
    result3 = agent_executor.invoke({"input": "如果再加一件200元的薄外套，总预算多少？"})
    print("Agent 回答：", result3["output"])